package cn.doitedu.ml.gender

import org.apache.log4j.{Level, Logger}
import org.apache.spark.ml.classification.{LogisticRegression, NaiveBayes}
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
import org.apache.spark.ml.linalg
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.mllib.evaluation.MulticlassMetrics
import org.apache.spark.mllib.linalg.Matrix
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.expressions.UserDefinedFunction

import scala.collection.mutable

object ActionGenderModelTrainner {

  def main(args: Array[String]): Unit = {

    Logger.getLogger("org.apache").setLevel(Level.WARN)

    val spark = SparkSession
      .builder()
      .appName("行为性别预测模型训练")
      .master("local")
      .getOrCreate()
    import org.apache.spark.sql.functions._
    import spark.implicits._

    val arr2Vec: UserDefinedFunction = udf((arr:mutable.WrappedArray[Double])=>{
      val vector: linalg.Vector = Vectors.dense(arr.toArray)
      vector
    })


    // 加载样本数据
    val sample = spark.read.option("header",true).option("inferSchema",true).csv("userprofile/data/gender/sample")

    /**
     * 特征工程，训练集向量化
     */
    // 用于训练朴素贝叶斯模型的特征
    val vec1 = sample.select('label,arr2Vec(array('category1,'category2,'category3,'brand1,'brand2,'brand3)) as "features")

    // 用于训练逻辑回归模型的特征
    val vec2 = sample.select('label,arr2Vec(array('day30_buy_cnts,'day30_buy_amt))  as "features")


    /**
     * 模型训练
     */
    // 构造朴素贝叶斯算法工具
    val naiveBayes = new NaiveBayes()
      .setLabelCol("label")
      .setFeaturesCol("features")
      .setSmoothing(1.0)
      .setModelType("multinomial")  // 多项式概率分布算法

    val bayesModel = naiveBayes.fit(vec1)


    // 构造逻辑回归算法工具
    val logisticRegression = new LogisticRegression()
      .setRegParam(0.1)
      .setFeaturesCol("features")
      .setLabelCol("label")

    val regrModel = logisticRegression.fit(vec2)

    /**
     * 加载测试数据来测试模型，评估
     */
    val test = spark.read.option("header",true).option("inferSchema",true).csv("userprofile/data/gender/test")
    val vecTest1 = test.select('guid,'label,arr2Vec(array('category1,'category2,'category3,'brand1,'brand2,'brand3)) as "features")
    val vecTest2 = test.select('guid,'label,arr2Vec(array('day30_buy_cnts,'day30_buy_amt)) as "features")

    val bayesPredict = bayesModel.transform(vecTest1)
    val regrPredict = regrModel.transform(vecTest2)

     bayesPredict.printSchema()
     bayesPredict.show(100,false)
     regrPredict.show(100,false)

    val vec2arr = (vec:linalg.Vector) =>{
      vec.toArray
    }
    spark.udf.register("vec2arr",vec2arr)


    // 最后，做一个加权综合，得到最终预测结果
    bayesPredict.createTempView("bayes")
    regrPredict.createTempView("regr")

    val result = spark.sql(
      """
        |select
        |bayes.guid,
        |bayes.label,
        |cast(if((vec2arr(bayes.probability)[0]*0.2 + vec2arr(regr.probability)[0]*0.8)<0.5,1.0,0.0) as double) as prediction
        |
        |from  bayes join regr on bayes.guid=regr.guid
        |
        |""".stripMargin)


    result.show(100,false)

    // 再对最终预测结果进行评估（混淆矩阵）
    val rdd = result.rdd.map(row=>{
      val label = row.getAs[Double]("label")
      val prediction = row.getAs[Double]("prediction")
      (prediction,label)
    })
    val matrix: Matrix = new MulticlassMetrics(rdd).confusionMatrix

    println(matrix)

    /**    P    N
     * P  1.0  1.0
     * N  0.0  2.0
     */

    spark.close()


  }

}
